import json
import os


def coco_to_yolo(coco_file, output_dir, image_dir):
    """
    此函数用于将COCO格式的标注文件转换为YOLO格式的标注文件。

    :param coco_file: coco格式的标注文件路径
    :param output_dir: 存储转换后的YOLO格式标注文件的目录路径。
    :param image_dir: 图像文件所在的目录路径。
    :return: None
    """
    try:
        # 读取COCO格式的标注文件
        with open(coco_file, 'r', encoding='utf-8') as f:
            coco_data = json.load(f)

        # 构建图像ID到图像信息的映射字典
        images = {img['id']: img for img in coco_data['images']}
        # 构建类别ID到类别名称的映射字典
        categories = {cat['id']: cat['name'] for cat in coco_data['categories']}
        # 构建类别ID到YOLO类别ID的映射字典
        category_id_to_yolo_id = {category_id: index for index, category_id in enumerate(categories.keys())}

        # 遍历COCO标注文件中的每个标注信息
        for annotation in coco_data['annotations']:
            image_id = annotation['image_id']
            image_info = images[image_id]
            image_width = image_info['width']
            image_height = image_info['height']

            category_id = annotation['category_id']
            # 获取对应的YOLO类别ID
            class_id = category_id_to_yolo_id[category_id]

            bbox = annotation['bbox']
            # 计算YOLO格式的边界框中心坐标
            x_center = (bbox[0] + bbox[2] / 2) / image_width
            y_center = (bbox[1] + bbox[3] / 2) / image_height
            # 计算YOLO格式的边界框宽度和高度
            width = bbox[2] / image_width
            height = bbox[3] / image_height

            image_filename = image_info['file_name']
            txt_filename = os.path.splitext(image_filename)[0] + '.txt'
            txt_path = os.path.join(output_dir, txt_filename)

            # 以追加模式打开文件并写入YOLO格式的标注信息
            with open(txt_path, 'a', encoding='utf-8') as f:
                f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")
    except FileNotFoundError:
        print(f"错误: 文件 {coco_file} 未找到。")
    except KeyError as e:
        print(f"错误: 在COCO数据中找不到键 {e}。")
    except Exception as e:
        print(f"发生未知错误: {e}")


if __name__ == "__main__":
    # 示例调用
    coco_file = '../datasets/annotations/modified_train.json'
    output_dir = '../datasets/labels/train'
    image_dir = '../datasets/images/train'
    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)
    coco_to_yolo(coco_file, output_dir, image_dir)
